import pygame
import sys

# Define grid dimensions (maze)
rows, cols = 3, 3

# Define graph-like connections in the maze (represent as node connections)
graph = {
    (0, 0): [(0, 1), (1, 0)],
    (0, 1): [(0, 0), (0, 2), (1, 1)],
    (0, 2): [(0, 1)],
    (1, 0): [(0, 0), (1, 1), (2, 0)],
    (1, 1): [(1, 0), (0, 1), (2, 1), (1, 2)],
    (1, 2): [(1, 1)],
    (2, 0): [(1, 0), (2, 1)],
    (2, 1): [(2, 0), (1, 1), (2, 2)],
    (2, 2): [(2, 1)]
}

# Define colors
RECT_COLOR = (255, 182, 193)  # Pink for cells
DOT_COLOR = (0, 0, 255)  # Blue for dots
TREE_EDGE_COLOR = (255, 0, 0)  # Red for tree edges

# Initialize pygame
pygame.init()

# Set screen size
WIDTH, HEIGHT = 600, 400
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Maze to Tree Transformation")

# Define positions for grid cells (manually placed for visualization)
cell_size = 80
offset_x, offset_y = 100, 50

# Define positions for the final tree structure
tree_positions = {}
tree_gap_x, tree_gap_y = 100, 80

# Start position in the final tree
root_position = (WIDTH // 2, 50)

# Define grid positions
node_positions = {}
for i in range(rows):
    for j in range(cols):
        node_positions[(i, j)] = (offset_x + j * cell_size, offset_y + i * cell_size)

# Set the font
font = pygame.font.SysFont(None, 24)

# Variables to control animation states
state = 0  # 0 = Show maze, 1 = Transform to dots, 2 = Build tree
dot_positions = {}  # Keep track of dot positions for animation
tree_edges = []  # Edges in the final tree

def draw_maze():
    """Draw the initial maze with rectangular cells."""
    for (i, j), (x, y) in node_positions.items():
        pygame.draw.rect(screen, RECT_COLOR, (x, y, cell_size, cell_size), 2)
        text = font.render(f"({i}, {j})", True, (0, 0, 0))
        screen.blit(text, (x + 10, y + 10))

def draw_dots():
    """Draw the nodes as dots, transforming from the maze."""
    for (i, j), pos in node_positions.items():
        if (i, j) in dot_positions:
            pygame.draw.circle(screen, DOT_COLOR, dot_positions[(i, j)], 10)
            text = font.render(f"{i},{j}", True, (0, 0, 0))
            screen.blit(text, (dot_positions[(i, j)][0] - 10, dot_positions[(i, j)][1] - 10))

def animate_to_dots():
    """Animate the transformation from rectangular cells to dots."""
    for node, (x, y) in node_positions.items():
        target_x, target_y = x + cell_size // 2, y + cell_size // 2
        if node not in dot_positions:
            dot_positions[node] = (x, y)
        current_x, current_y = dot_positions[node]
        if abs(current_x - target_x) > 2 or abs(current_y - target_y) > 2:
            new_x = current_x + (target_x - current_x) // 10
            new_y = current_y + (target_y - current_y) // 10
            dot_positions[node] = (new_x, new_y)
        else:
            dot_positions[node] = (target_x, target_y)

def build_tree():
    """Build the tree by re-arranging the dots into a tree structure."""
    # Calculate the tree positions for each node
    def arrange_tree(node, level, position):
        if node not in tree_positions:
            tree_positions[node] = position
        children = [n for n in graph[node] if n not in tree_positions]
        n = len(children)
        for i, child in enumerate(children):
            x_offset = (i - (n - 1) / 2) * tree_gap_x
            child_pos = (position[0] + x_offset, position[1] + tree_gap_y)
            tree_edges.append((node, child))
            arrange_tree(child, level + 1, child_pos)

    if not tree_positions:
        arrange_tree((0, 0), 0, root_position)

    # Animate the movement of dots to their final tree positions
    for node, target_pos in tree_positions.items():
        if node not in dot_positions:
            continue
        current_pos = dot_positions[node]
        if abs(current_pos[0] - target_pos[0]) > 2 or abs(current_pos[1] - target_pos[1]) > 2:
            new_x = current_pos[0] + (target_pos[0] - current_pos[0]) // 10
            new_y = current_pos[1] + (target_pos[1] - current_pos[1]) // 10
            dot_positions[node] = (new_x, new_y)
        else:
            dot_positions[node] = target_pos

    # Draw tree edges
    for edge in tree_edges:
        node1, node2 = edge
        pygame.draw.line(screen, TREE_EDGE_COLOR, dot_positions[node1], dot_positions[node2], 2)

def main():
    global state
    clock = pygame.time.Clock()
    running = True

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    state += 1  # Move to the next state

        # Clear the screen
        screen.fill((255, 255, 255))

        # Draw the maze, then animate transformation to dots, then build the tree
        if state == 0:
            draw_maze()  # Initial maze with rectangular cells
        elif state == 1:
            animate_to_dots()  # Transform the cells into dots
            draw_dots()
        elif state == 2:
            build_tree()  # Build the tree from the dots
            draw_dots()

        pygame.display.flip()
        clock.tick(30)

    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()